package com.threadpool.pool;

import com.threadpool.pool.policy.RejectPolicy;
import com.threadpool.pool.queue.CustomBlockingQueue;

import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;

/**
 * @author lizhentao
 */
public class CustomThreadPool<T extends Runnable> {

    private final int coreSize;
    private final int maxSize;
    private final long keepAliveTime;
    private final TimeUnit timeUnit;
    private final CustomBlockingQueue<T> taskQueue;
    private final RejectPolicy<T> rejectPolicy;
    private final CustomThreadFactory threadFactory;
    private final Set<Worker> workers = new HashSet<>();

    public CustomThreadPool(int coreSize, int maxSize, long keepAliveTime, TimeUnit timeUnit,
                            CustomBlockingQueue<T> taskQueue,
                            CustomThreadFactory threadFactory,
                            RejectPolicy<T> rejectPolicy) {
        this.coreSize = coreSize;
        this.maxSize = maxSize;
        this.keepAliveTime = keepAliveTime;
        this.timeUnit = timeUnit;
        this.taskQueue = taskQueue;
        this.threadFactory = threadFactory;
        this.rejectPolicy = rejectPolicy;
    }

    public CustomBlockingQueue<T> getTaskQueue() {
        return taskQueue;
    }

    public void execute(T task) {
        synchronized (workers) {
            if (workers.size() < coreSize) {
                Worker worker = new Worker(task);
                Thread thread = threadFactory.newThread(worker);
                workers.add(worker);
                thread.start();
                return;
            }

            boolean offered = taskQueue.offer(task);
            if (offered) {
                return;
            }

            // 如果队列满了但线程数还没到最大值，可以再扩充线程
            if (workers.size() < maxSize) {
                Worker worker = new Worker(task);
                Thread thread = threadFactory.newThread(worker);
                workers.add(worker);
                thread.start();
            } else {
                // 超过最大线程数，执行拒绝策略
                rejectPolicy.reject(task, this);
            }
        }
    }

    private class Worker implements Runnable {
        private T firstTask;

        public Worker(T task) {
            this.firstTask = task;
        }

        @Override
        public void run() {
            T task = this.firstTask;
            while (task != null || (task = taskQueue.poll(keepAliveTime, timeUnit)) != null) {
                try {
                    task.run();
                } finally {
                    task = null;
                }
            }
            synchronized (workers) {
                workers.remove(this);
            }
        }
    }
}
